'''
Created on Jul 13, 2012

@author: nzhao
'''
import numpy as np
from traits.api import HasTraits, Array, Int, Enum, Float
from phy.general.Constant import SpinDictionary
from math_qm.Constant import X, Y, Z, SPIN1X, SPIN1Y, SPIN1Z
 
class Hamiltonian(HasTraits):
    matrix     = Array()
    spin_num   = Int()
    spin_value = Float()
    spin_zfs   = Float(0.0)
    dim        = Int()
    
    eig_vals   = Array()
    eig_vect   = Array()
    
    def calculate_eigs(self):
        vals, vect = np.linalg.eigh(self.matrix)
        idx = vals.argsort()
        self.eig_vals = vals[idx]
        self.eig_vect = vect[:,idx]
        

class SingleSpin_Hamiltonian(Hamiltonian):    
    spin_element = Enum('NV', 'E', 'C13', 'N14', 'N15', 'H')
    
    def __init__(self, spin_element='NV'):
        self.spin_num     = 1
        self.spin_element = spin_element
        self.spin_value   = SpinDictionary[spin_element]['SpinValue']
        self.dim          = int(2 * self.spin_value + 1)
        self.matrix       = np.zeros( (self.dim, self.dim) )

        if self.spin_value > 0.5:
            self.zfs = SpinDictionary[self.spin_element]['ZeroFieldSplitting']
    
    def calculate_hamiltonian(self,  magnetic_field):
        bx, by, bz = magnetic_field.xyz * SpinDictionary[self.spin_element]['GyromagneticRatio']
        
        if self.spin_value == 0.5:
            self.matrix = 0.5 * ( bx * X + by * Y + bz * Z) / 1000.0 #k rad/s 
        elif self.spin_value == 1.0:
            self.matrix = ( bx * SPIN1X + by * SPIN1Y + bz * SPIN1Z ) / 1000.0 #k rad/s 
            self.matrix = self.matrix + self.zfs * np.dot(SPIN1Z, SPIN1Z)
        else:
            self.matrix = np.zeros( (self.dim, self.dim) )
         

class MultiSpin_Hamiltonian(Hamiltonian):
    pass
